from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, insert
from sqlalchemy.orm import selectinload
from typing import List, Tuple
import math
from app.crud.base import CRUDBase
from app.models.chat import Conversation, Message, message_files
from app.models.file import UploadedFile
from app.schemas.chat import ConversationCreate, MessageCreate


class CRUDConversation(CRUDBase[Conversation, ConversationCreate, None]):
    async def get_by_user(self, db: AsyncSession, *, user_id: UUID) -> list[Conversation]:
        result = await db.execute(select(self.model).filter(self.model.user_id == user_id))
        return result.scalars().all()
    
    async def get_user_conversations_paginated(
        self, 
        db: AsyncSession, 
        *, 
        user_id: UUID, 
        page: int = 1, 
        page_size: int = 20
    ) -> Tuple[List[Conversation], int]:
        """获取用户对话列表（分页）"""
        # Calculate offset
        offset = (page - 1) * page_size
        
        # Get total count
        count_result = await db.execute(
            select(func.count(self.model.id)).filter(self.model.user_id == user_id)
        )
        total = count_result.scalar() or 0
        
        # Get paginated results
        conversations_result = await db.execute(
            select(self.model)
            .filter(self.model.user_id == user_id)
            .order_by(self.model.updated_at.desc())
            .offset(offset)
            .limit(page_size)
        )
        conversations = conversations_result.scalars().all()
        
        return list(conversations), total
    
    async def get_conversation_message_counts(
        self, 
        db: AsyncSession, 
        *, 
        conversation_ids: List[UUID]
    ) -> dict[UUID, int]:
        """获取多个对话的消息数量"""
        if not conversation_ids:
            return {}
            
        result = await db.execute(
            select(Message.conversation_id, func.count(Message.id).label('count'))
            .filter(Message.conversation_id.in_(conversation_ids))
            .group_by(Message.conversation_id)
        )
        
        # 转换为字典格式
        counts = {}
        for row in result:
            counts[row.conversation_id] = row.count
        
        # 为没有消息的对话设置为0
        for conv_id in conversation_ids:
            if conv_id not in counts:
                counts[conv_id] = 0
                
        return counts

    async def get_conversations_by_file_id(self, db: AsyncSession, *, file_id: UUID) -> List[UUID]:
        """根据文件ID获取关联的对话ID列表"""
        result = await db.execute(
            select(Message.conversation_id)
            .join(message_files, Message.id == message_files.c.message_id)
            .filter(message_files.c.file_id == file_id)
            .distinct()
        )
        return [row[0] for row in result.fetchall()]


class CRUDMessage(CRUDBase[Message, MessageCreate, None]):
    async def get_by_conversation(self, db: AsyncSession, *, conversation_id: UUID) -> list[Message]:
        result = await db.execute(
            select(self.model)
            .options(selectinload(self.model.files))
            .filter(self.model.conversation_id == conversation_id)
            .order_by(self.model.created_at.asc())
        )
        return result.scalars().all()
    
    async def get_by_conversation_paginated(
        self, 
        db: AsyncSession, 
        *, 
        conversation_id: UUID, 
        page: int = 1, 
        page_size: int = 20
    ) -> Tuple[List[Message], int]:
        """获取对话消息列表（分页）"""
        # Calculate offset
        offset = (page - 1) * page_size
        
        # Get total count
        count_result = await db.execute(
            select(func.count(self.model.id))
            .filter(self.model.conversation_id == conversation_id)
        )
        total = count_result.scalar() or 0
        
        # Get paginated results with files (按时间升序，最早的消息在前)
        messages_result = await db.execute(
            select(self.model)
            .options(selectinload(self.model.files))
            .filter(self.model.conversation_id == conversation_id)
            .order_by(self.model.created_at.asc())
            .offset(offset)
            .limit(page_size)
        )
        messages = messages_result.scalars().all()
        
        return list(messages), total

    async def create(self, db: AsyncSession, *, obj_in: MessageCreate) -> Message:
        # Create the message
        db_obj = self.model(
            conversation_id=obj_in.conversation_id,
            role=obj_in.role,
            content=obj_in.content,
            reasoning_content=obj_in.reasoning_content,
        )
        db.add(db_obj)
        await db.flush()
        
        # Add file associations if provided
        if obj_in.file_ids:
            for file_id in obj_in.file_ids:
                # Verify file exists and belongs to the same user (optional validation)
                file_obj = await db.execute(select(UploadedFile).filter(UploadedFile.id == file_id))
                file_record = file_obj.scalar_one_or_none()
                if file_record:
                    # Insert into message_files association table directly
                    await db.execute(
                        insert(message_files).values(
                            message_id=db_obj.id,
                            file_id=file_id
                        )
                    )
        
        await db.commit()
        await db.refresh(db_obj)
        return db_obj


crud_conversation = CRUDConversation(Conversation)
crud_message = CRUDMessage(Message)
